#include "stdafx.h"
#include "kernels_cpu.h"

#ifndef Epsilon
#define Epsilon 1e-4f
#endif

namespace Kernels
{
struct uint2
{
	unsigned int x,y;
};

struct float2
{
	float x,y;
	float2 operator * (float f) const
	{
		float2 rs;
		rs.x = x * f;
		rs.y = y * f;
		return rs;
	}
	float2 operator + (float2 v) const
	{
		float2 rs;
		rs.x = x + v.x;
		rs.y = y + v.y;
		return rs;
	}
};

union intAsFloat
{
	int vint;
	float vfloat;
};

float as_float(unsigned int v)
{
	intAsFloat c;
	c.vint = v;
	return c.vfloat;
}

struct float3
{
	float x,y,z; float pading;

	float3 operator * (float f) const
	{
		float3 rs;
		rs.x = x * f;
		rs.y = y * f;
		rs.z = z * f;
		return rs;
	}
	float3 operator + (float3 v) const
	{
		float3 rs;
		rs.x = x + v.x;
		rs.y = y + v.y;
		rs.z = z + v.z;
		return rs;
	}
};
	
float3 make_float3(float v1, float v2, float v3)
{
	float3 rs;
	rs.x = v1;
	rs.y = v2;
	rs.z = v3;
	return rs;
}
float dot(float3 v1, float3 v2)
{
	return v1.x * v2.x + v1.y * v2.y + v1.z * v2.z;
}

struct float4
{
	union
	{
		struct
		{
			float x, y, z, w;
		};
		float3 xyz;
	};

	float4 operator * (float f) const
	{
		float4 rs;
		rs.x = x * f;
		rs.y = y * f;
		rs.z = z * f;
		rs.w = w * f;
		return rs;
	}
	float4 operator + (float4 v) const
	{
		float4 rs;
		rs.x = x + v.x;
		rs.y = y + v.y;
		rs.z = z + v.z;
		rs.w = w + v.w;
		return rs;
	}
};

struct Vec3
{
	union
	{
		struct
		{
			float x,y,z;
		};
		float v[3];
	};

};

struct Vec3 MakeVec3(float3 v)
{
	struct Vec3 rs;
	rs.x = v.x;
	rs.y = v.y;
	rs.z = v.z;
	return rs;
}

struct Vec3 MakeVec3f(float x, float y, float z)
{
	struct Vec3 rs;
	rs.x = x;
	rs.y = y;
	rs.z = z;
	return rs;
}

float Dot(const struct Vec3 v1, const struct Vec3 v2)
{
	return v1.x*v2.x + v1.y*v2.y + v1.z*v2.z;
}

float Dot_float3(const struct Vec3 v1, const float3 v2)
{
	return v1.x*v2.x + v1.y*v2.y + v1.z*v2.z;
}

struct matrix4
{
	float4 c1, c2, c3, c4;
};

float3 mul3(const struct matrix4 mat, const float3 vec)
{
	return mat.c1.xyz * vec.x + mat.c2.xyz * vec.y + mat.c3.xyz * vec.z + mat.c4.xyz;
}

float4 mul(const struct matrix4 mat, const float4 vec)
{
	return mat.c1 * vec.x + mat.c2 * vec.y + mat.c3 * vec.z + mat.c4 * vec.w;
}

float3 transformNormal(const struct matrix4 mat, const float3 vec)
{
	return mat.c1.xyz * vec.x + mat.c2.xyz * vec.y + mat.c3.xyz * vec.z;
}

struct Ray
{
	float3 Origin;
	float3 Direction;
	float3 ReciprocalDirection;
	float tMin, tMax; 
	float pading1, pading2;
};

struct RayDifferential
{
	float3 dPdx, dPdy, dDdx, dDdy;
};

float Take(float3 v, int c)
{
	return ((float*)&v)[c];
}

void Set(float3 * v, int i, float val)
{
	((float*)v)[i] = val;
}

/*kernel*/
void GenerateCameraRays (
					int global_id_0, int global_id_1,
					const int width, 
					const int height,
					const float4 right4, 
					const float4 up4,
					const float centerX, 
					const float centerY, 
					const float screenZ,
					const struct matrix4 cameraTransform,
                     /*__global*/ struct Ray* rays,
                     /*__global*/ struct RayDifferential * rayDiffs)
{
	float3 right = right4.xyz;
	float3 up = up4.xyz;
	int i = global_id_0;
	int j = global_id_1;
	if (i >= width || j >= height)
		return;
	struct Ray r;
	struct RayDifferential rayDiff;
	// Compute camera ray
	// dir = || right * px + up * py + screenZ ||
	r.tMin = Epsilon;
	r.tMax = FLT_MAX;
	float3 d = right * ((float)i-centerX) + up * (centerY-(float)j) + make_float3(0.0f, 0.0f, screenZ);
	float dirDOTdir = dot(d, d);
	float dirLength = sqrt(dirDOTdir);
	float invDirLength = 1.0f/dirLength;
	r.Direction = d * invDirLength;
	r.Origin = mul3(cameraTransform, make_float3(0.0f, 0.0f, 0.0f));
	r.Direction = transformNormal(cameraTransform, r.Direction);
	r.ReciprocalDirection.x = 1.0f / r.Direction.x;
	r.ReciprocalDirection.y = 1.0f / r.Direction.y;
	r.ReciprocalDirection.z = 1.0f / r.Direction.z;

	// Compute initial ray differential
	rayDiff.dPdx = make_float3(0.0f, 0.0f, 0.0f);
	rayDiff.dPdy = make_float3(0.0f, 0.0f, 0.0f);
		
	float dirDOTdir3o2 = 1.0f / sqrt(dirDOTdir * dirDOTdir * dirDOTdir);
	rayDiff.dDdx = right * (dirDOTdir - dot(d, right)) * dirDOTdir3o2;
	rayDiff.dDdy = up * (dirDOTdir-dot(d, up)) * dirDOTdir3o2;

	rays[i*height + j] = r;
	rayDiffs[i*height + j] = rayDiff;
}


////////////////// Ray Intersection //////////////////

struct MeshFace
{
	float PlaneU, PlaneV, PlaneD;
	short MaterialId;
	short ProjectionAxis;
	float K_beta_u, K_beta_v, K_beta_d, K_alpha_u;
	float K_gamma_u, K_gamma_v, K_gamma_d, K_alpha_v;
	union
	{
		struct
		{
			int Normal1, Normal2, Normal3;
		};
		int Normals[3];
	};
	union
	{
		struct
		{
			int TexCoord1, TexCoord2, TexCoord3;
		};
		int TexCoords[3];
	};
};

struct DifferentialGeometry
{
	struct Vec3 Position;
	struct Vec3 Normal;
	float2 UV;
	int MaterialId;
	float dUdx, dUdy, dVdx, dVdy;
	int dNdxContinuation; // faceId if >=0, -1 if computed
	struct Vec3 dNdx, dNdy; // Lazy computation
};

struct BBox
{
	struct Vec3 Min;
	struct Vec3 Max;
};

struct KdLeafHeader
{
	struct BBox Bounds;
	int Ropes[6];
};

int RayBBoxIntersection(const struct BBox bbox, const struct Ray ray, float * tmin, float * tmax)
{
	float tymin, tymax, tzmin, tzmax;
	if (ray.ReciprocalDirection.x >= 0) 
	{
		*tmin = (bbox.Min.x - ray.Origin.x) * ray.ReciprocalDirection.x;
		*tmax = (bbox.Max.x - ray.Origin.x) * ray.ReciprocalDirection.x;
	}
	else
	{
		*tmin = (bbox.Max.x - ray.Origin.x) * ray.ReciprocalDirection.x;
		*tmax = (bbox.Min.x - ray.Origin.x) * ray.ReciprocalDirection.x;
	}
	if (ray.ReciprocalDirection.y >= 0) 
	{
		tymin = (bbox.Min.y - ray.Origin.y) * ray.ReciprocalDirection.y;
		tymax = (bbox.Max.y - ray.Origin.y) * ray.ReciprocalDirection.y;
	}
	else
	{
		tymin = (bbox.Max.y - ray.Origin.y) * ray.ReciprocalDirection.y;
		tymax = (bbox.Min.y - ray.Origin.y) * ray.ReciprocalDirection.y;
	}
	if ( *tmin - tymax > Epsilon || tymin - *tmax > Epsilon)
		return 0;
	if (tymin > *tmin)
		*tmin = tymin;
	if (tymax < *tmax)
		*tmax = tymax;
	if (ray.ReciprocalDirection.z >= 0) 
	{
		tzmin = (bbox.Min.z - ray.Origin.z) * ray.ReciprocalDirection.z;
		tzmax = (bbox.Max.z - ray.Origin.z) * ray.ReciprocalDirection.z;
	}
	else
	{
		tzmin = (bbox.Max.z - ray.Origin.z) * ray.ReciprocalDirection.z;
		tzmax = (bbox.Min.z - ray.Origin.z) * ray.ReciprocalDirection.z;
	}
	if ( *tmin - tzmax > Epsilon || tzmin - *tmax > Epsilon )
		return 0;
	if (tzmin > *tmin)
		*tmin = tzmin;
	if (tzmax < *tmax)
		*tmax = tzmax;
	*tmin = max(ray.tMin, *tmin);
	*tmax = min(ray.tMax, *tmax);
	return *tmin < *tmax;
}

int NodeIsLeaf(uint2 node)
{
	return (int)((node.x) & (1<<31));
}

int NodeGetSplitDimension(uint2 node)
{
	return (int)((node.x) & 3);
}

float NodeGetSplitPlane(uint2 node)
{
	return as_float(node.y);
}

int NodeGetElementCount(uint2 node)
{
	return (int)(node.x & 0x7FFFFFFF);
}

int NodeGetOffset(uint2 node)
{
	return (int)(node.x & 0x7FFFFFFC);
}

uint2 FetchNode(const uint2 * kdNodes, int byteOffset)
{
	return *(const uint2*)((const char*)kdNodes + byteOffset);
}

int RayBBoxExit(const struct BBox bbox, const struct Ray ray, float tmin, float * tmax)
{
	int face;
	float tymax, tzmax;
	int bufFace;
	if (ray.ReciprocalDirection.x >= 0) 
	{
		*tmax = (bbox.Max.x - ray.Origin.x) * ray.ReciprocalDirection.x;
		face = 3;
	}
	else
	{
		*tmax = (bbox.Min.x - ray.Origin.x) * ray.ReciprocalDirection.x;
		face = 0;
	}
	if (ray.ReciprocalDirection.y >= 0) 
	{
		tymax = (bbox.Max.y - ray.Origin.y) * ray.ReciprocalDirection.y;
		bufFace = 4;
	}
	else
	{
		tymax = (bbox.Min.y - ray.Origin.y) * ray.ReciprocalDirection.y;
		bufFace = 1;
	}
	if (tymax > tmin && tymax < *tmax)
	{
		*tmax = tymax;
		face = bufFace;
	}
	if (ray.ReciprocalDirection.z >= 0) 
	{
		tzmax = (bbox.Max.z - ray.Origin.z) * ray.ReciprocalDirection.z;
		bufFace = 5;
	}
	else
	{
		tzmax = (bbox.Min.z - ray.Origin.z) * ray.ReciprocalDirection.z;
		bufFace = 2;
	}
	if (tzmax > tmin && tzmax < *tmax)
	{
		*tmax = tzmax;
		face = bufFace;
	}
	*tmax = min(ray.tMax, *tmax);
	return face;
}

int TraceFace(struct DifferentialGeometry * inter, const struct MeshFace face, int faceId, 
	const struct Vec3 * normals, 
	const float2 * texCoords, const struct Ray ray, const struct RayDifferential diffIn, float * t, int id)
{
	const int mod3[] = {0,1,2,0,1};
	int u = mod3[face.ProjectionAxis + 1];
	int v = mod3[face.ProjectionAxis + 2];
	float originU = Take(ray.Origin, u);
	float originV = Take(ray.Origin, v);
	float originW = Take(ray.Origin, face.ProjectionAxis);
	float dirU = Take(ray.Direction, u);
	float dirV = Take(ray.Direction, v);
	float dirW = Take(ray.Direction, face.ProjectionAxis);
	float invNdotD = 1.0f / ( face.PlaneU * dirU
							  + face.PlaneV * dirV 
					          + dirW );
	float tplane = -(face.PlaneU * originU 
					+ face.PlaneV * originV 
					+ originW + face.PlaneD) * invNdotD;
	if (id == 0 && faceId == 15392)
	{
		printf("%d, %d\n", u, v);
		printf("%f %f %f %f %f %f\n", originU, originV, originW, dirU, dirV, dirW);
		printf("%f %f %f %f", invNdotD, tplane, ray.tMin, ray.tMax);
	}
	if (tplane < ray.tMin/*-Epsilon*/ || tplane > ray.tMax/*+Epsilon*/)
		return 0;
	float hitU = originU + dirU * tplane;
	float hitV = originV + dirV * tplane;
			
	float beta = face.K_beta_u * hitU + face.K_beta_v * hitV + face.K_beta_d;
	if (beta < -Epsilon)
		return 0;
	float gamma = face.K_gamma_u * hitU + face.K_gamma_v * hitV + face.K_gamma_d;
	if (gamma < -Epsilon)
		return 0;
	if (beta + gamma > 1.0f + Epsilon)
		return 0;
	// if (pred)
	//	 return true;
	inter->Position.v[u] = hitU;
	inter->Position.v[v] = hitV;
	inter->Position.v[face.ProjectionAxis] = originW + dirW * tplane;
	float alpha = 1.0f - beta - gamma;

	struct Vec3 norm1 = normals[face.Normal1];
	struct Vec3 norm2 = normals[face.Normal2];
	struct Vec3 norm3 = normals[face.Normal3];
	inter->Normal.x = norm1.x * alpha + norm2.x * beta + norm3.x * gamma;
	inter->Normal.y = norm1.y * alpha + norm2.y * beta + norm3.y * gamma;
	inter->Normal.z = norm1.z * alpha + norm2.z * beta + norm3.z * gamma;
	if (face.TexCoord1 != -1)
	{
		inter->UV = texCoords[face.TexCoord1] * alpha + texCoords[face.TexCoord2] * beta + texCoords[face.TexCoord3] * gamma;
	}
	inter->dNdxContinuation = faceId;
	*t = tplane;
	inter->MaterialId = face.MaterialId;
	return 1;
}

/*__kernel*/
void RayMeshIntersection (			int global_id_0,
									int size,
									float GeometryBias,
									struct BBox bounds,
									int autoFlipNormals, 
									const struct Ray* rays,
									const struct RayDifferential * rayDiffs,
									const uint2 * kdNodes,
									const struct MeshFace* faces,
									const char* leaves,
									const struct Vec3 * normals,
									const float2 * texCoords,
									struct DifferentialGeometry * results,
									float* intersectionTs)
{
	int id = global_id_0;
	if (id >= size)
		return;
	struct Ray ray = rays[id];
	struct RayDifferential rayDiff = rayDiffs[id];
	struct DifferentialGeometry rs;
	float sceneMin, sceneMax;

	if (!RayBBoxIntersection(bounds, ray, &sceneMin, &sceneMax))
	{
		//results[id].dNdxContinuation = -2;
		intersectionTs[id] = FLT_MAX;
		return;
	}
	float t = FLT_MAX;
	struct Ray traceRay = ray;
	traceRay.tMax = sceneMax;
	traceRay.tMin = sceneMin;
	rs.Normal = MakeVec3(traceRay.Direction);
	bool intersected = false;
	float tMin = sceneMin;
	float tMax = sceneMax;
	int nodeId = 0;
	uint2 node = kdNodes[0];
	while (tMin < tMax)
	{
		while (!NodeIsLeaf(node))
		{
			int axis = NodeGetSplitDimension(node);
			float tplane;
			float splitPlane = NodeGetSplitPlane(node);
			float rayOrigin = ((float*)&ray.Origin)[axis];
			tplane = (splitPlane - rayOrigin) * ((float*)&ray.ReciprocalDirection)[axis];
			int node1, node2;
			if (rayOrigin < splitPlane)
			{
				node1 = nodeId + 8;
				node2 = NodeGetOffset(node);
			}
			else
			{
				node1 = NodeGetOffset(node);
				node2 = nodeId + 8;
			}
			if (tplane >= tMax || tplane < 0)
			{
				nodeId = node1;
			}
			else if (tplane <= tMin)
			{
				nodeId = node2;
			}
			else
			{
				nodeId = node1;
			}
			node = FetchNode(kdNodes, nodeId);
			
		}
		
		struct DifferentialGeometry inter;
		int count = NodeGetElementCount(node);
		const struct KdLeafHeader* leafPtr = (const struct KdLeafHeader*)((const char*)leaves + node.y);
		struct KdLeafHeader leaf = *leafPtr;
		const int* elemIds = (const int*)(leafPtr+1);
		
		for (int i = 0; i < count; i++)
		{
			int faceId = elemIds[i];
			struct MeshFace curElem = faces[faceId];
			float it;
			
			if (TraceFace(&inter, curElem, faceId, normals, texCoords, traceRay, rayDiff, &it, id))
			{
				if (id == 0)
				{
					printf("%d-%d-%f\n ", nodeId, faceId, it);
	
				}
				// bias t
				it += GeometryBias*Dot_float3(inter.Normal, ray.Direction);
				if (it<tMax)
				{
					tMax = it;
					rs = inter;
					t = it;
				}
			}
		}
		// Locate next node
		int face;
		float tMin_cpy = tMin;
		face = RayBBoxExit(leaf.Bounds, ray, tMin_cpy, &tMin);
		tMin -= GeometryBias;
		int nextPtr = leaf.Ropes[face];
		if (nextPtr == -1)
			break;
		nodeId = nextPtr;
		node = FetchNode(kdNodes, nodeId);
	}
	intersectionTs[id] = t;
	float len = sqrt(rs.Normal.x*rs.Normal.x + rs.Normal.y*rs.Normal.y + rs.Normal.z*rs.Normal.z);
	rs.Normal.x /= len;
	rs.Normal.y /= len;
	rs.Normal.z /= len;
	if (autoFlipNormals && t!= FLT_MAX && Dot_float3(rs.Normal, ray.Direction) > 0.0f)
	{
		rs.Normal.x = -rs.Normal.x;
		rs.Normal.y = -rs.Normal.y;
		rs.Normal.z = -rs.Normal.z;
	}
	results[id] = rs;
}


void RayMeshIntersectionCPUKernel (			int global_id_0,
									int size,
									float GeometryBias,
									RayTracePro::BBox bounds,
									int autoFlipNormals, 
									const void* rays,
									const void * rayDiffs,
									const void * kdNodes,
									const void * faces,
									const void * leaves,
									const void * normals,
									const void * texCoords,
									void * results,
									float* intersectionTs)
{
	Kernels::BBox box;
	memcpy(&box, &bounds, sizeof(BBox));
	RayMeshIntersection(global_id_0, size, GeometryBias, box, autoFlipNormals, (const Ray*)rays,
		(const RayDifferential*) rayDiffs, (const uint2*)kdNodes, (const MeshFace*)faces,
		(const char*)leaves, (const Vec3*)normals, (const float2*)texCoords, (DifferentialGeometry*)results, intersectionTs);
}

}